{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Distributed Data Classification with NeMo Curator's `FineWebNemotronEduClassifier`\n",
    "\n",
    "This notebook demonstrates the use of NeMo Curator's `FineWebNemotronEduClassifier`. The [FineWeb Nemotron-4 Edu classifier](https://huggingface.co/nvidia/nemocurator-fineweb-nemotron-4-edu-classifier) is used to determine the educational value (score 0-5 from low to high) of a text. It helps with data annotation, which is useful in data blending for foundation model training. Please refer to the Hugging Face page for more information about the NemoCurator FineWeb Nemotron-4 Edu Classifier, including its output labels, here: https://huggingface.co/nvidia/nemocurator-fineweb-nemotron-4-edu-classifier.\n",
    "\n",
    "This tutorial requires at least 1 NVIDIA GPU with:\n",
    "  - Volta™ or higher (compute capability 7.0+)\n",
    "  - CUDA 12.x\n",
    "\n",
    "Before running this notebook, see this [Installation Guide](https://docs.nvidia.com/nemo/curator/latest/admin/installation.html#admin-installation) page for instructions on how to install NeMo Curator. Be sure to use an installation method which includes GPU dependencies.\n",
    "\n",
    "For more information about the classifiers, refer to our [Distributed Data Classification](https://docs.nvidia.com/nemo/curator/latest/curate-text/process-data/quality-assessment/distributed-classifier.html) documentation page."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Silence Curator logs via Loguru\n",
    "import os\n",
    "\n",
    "os.environ[\"LOGURU_LEVEL\"] = \"ERROR\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following imports are required for this tutorial:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "from nemo_curator.core.client import RayClient\n",
    "from nemo_curator.pipeline import Pipeline\n",
    "from nemo_curator.stages.text.classifiers import FineWebNemotronEduClassifier\n",
    "from nemo_curator.stages.text.io.reader.jsonl import JsonlReader\n",
    "from nemo_curator.stages.text.io.writer.jsonl import JsonlWriter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run a pipeline in NeMo Curator, we must start a Ray cluster. This can be done manually (see the [Ray documentation](https://docs.ray.io/en/latest/ray-core/starting-ray.html)) or with Curator's `RayClient`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    ray_client = RayClient()\n",
    "    ray_client.start()\n",
    "except Exception as e:\n",
    "    msg = f\"Error initializing Ray client: {e}\"\n",
    "    raise RuntimeError(msg) from e"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initialize Read, Classification, and Write Stages\n",
    "\n",
    "Functions in NeMo Curator are called stages. For this tutorial, we will initialize 3 stages: a JSONL file reader, a FineWeb Nemotron-4 Edu classification stage, and a JSONL file writer.\n",
    "\n",
    "For this tutorial, we will create a sample JSONL file to use:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_file_path = \"./input_data_dir\"\n",
    "\n",
    "# Create sample dataset for the tutorial\n",
    "text = [\n",
    "    \"Get rich quick by following these 5 simple steps!!! No experience needed, just click here now!!\",\n",
    "    \"Our website uses cookies to improve your experience. By continuing, you agree to our privacy policy.\",\n",
    "    \"The water cycle is a continuous process that moves water throughout the Earth's atmosphere, surface, and underground reservoirs. It includes key stages such as evaporation, where water from oceans, lakes, and rivers turns into vapor; condensation, where that vapor cools and forms clouds; and precipitation, where water returns to the surface as rain, snow, or hail. This cycle is essential because it regulates climate, replenishes freshwater supplies, and ensures that ecosystems and human communities have access to clean water.\",\n",
    "    \"George Washington was the first president of the United States, serving from 1789 to 1797 after playing a crucial role as commander of the Continental Army during the American Revolutionary War. He was born in 1732 and died in 1799.\",\n",
    "    \"Newton's three laws of motion form the foundation of classical mechanics. The first law states that an object will remain at rest or in uniform motion unless acted upon by an external force. The second law defines the relationship between force, mass, and acceleration (F = ma). The third law asserts that every action has an equal and opposite reaction.\",\n",
    "]\n",
    "df = pd.DataFrame({\"text\": text})\n",
    "\n",
    "try:\n",
    "    os.makedirs(input_file_path, exist_ok=True)\n",
    "    df.to_json(input_file_path + \"/data.jsonl\", orient=\"records\", lines=True)\n",
    "except Exception as e:\n",
    "    msg = f\"Error creating input file: {e}\"\n",
    "    raise RuntimeError(msg) from e"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can define the reader stage with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read existing directory of JSONL files\n",
    "read_stage = JsonlReader(input_file_path, files_per_partition=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The classifier stage is broken down under the hood into a tokenizer stage and a model inference stage. Tokenization is run on the CPU while model inference is run on the GPU. This means that behind the scenes, the `FineWebNemotronEduClassifier` stage is actually being broken down into 2 stages (some parameters and details omitted to avoid complexity, please refer to the documentation for more details):\n",
    "\n",
    "```python\n",
    "class TokenizerStage:\n",
    "    self.resources = Resources(cpus=1)\n",
    "    self.model_identifier = \"nvidia/nemocurator-fineweb-nemotron-4-edu-classifier\"\n",
    "    self.text_field = \"text\"\n",
    "    self.padding_side = \"right\"\n",
    "    ...\n",
    "class ModelStage:\n",
    "    self.resources = Resources(cpus=1, gpus=1)\n",
    "    self.model_identifier = \"nvidia/nemocurator-fineweb-nemotron-4-edu-classifier\"\n",
    "    self.model_inference_batch_size = 1024\n",
    "    ...\n",
    "```\n",
    "\n",
    "Optionally, the classifier predictions may be filtered to include only texts with values listed in `filter_by`. If the `filter_by` parameter is set, then a third stage is added:\n",
    "\n",
    "```python\n",
    "def filter_by_category(self, value: str) -> bool:\n",
    "    return value in self.filter_by\n",
    "\n",
    "...\n",
    "\n",
    "if self.filter_by is not None and len(self.filter_by) > 0:\n",
    "    self.stages.append(Filter(filter_fn=self.filter_by_category, filter_field=...))\n",
    "```\n",
    "\n",
    "Since the NemoCurator FineWeb Nemotron-4 Edu classifier outputs a floating point value ranging from 0 to 5, Curator labels samples with a value of 2.5 or higher as `high_quality` and samples with a value lower than 2.5 as `low_quality`.\n",
    "\n",
    "See the [API Reference](https://docs.nvidia.com/nemo/curator/latest/apidocs/stages/stages.text.classifiers.fineweb_edu.html#stages.text.classifiers.fineweb_edu.FineWebNemotronEduClassifier) for more information about the `FineWebNemotronEduClassifier` class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize the FineWeb Nemotron-4 Edu classifier\n",
    "classifier_stage = FineWebNemotronEduClassifier()\n",
    "\n",
    "# If desired, you may filter your dataset with:\n",
    "# classifier_stage = FineWebNemotronEduClassifier(filter_by=[\"high_quality\"])  # noqa: ERA001\n",
    "# or\n",
    "# classifier_stage = FineWebNemotronEduClassifier(filter_by=[\"low_quality\"])  # noqa: ERA001"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can define a stage for writing the results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Write results to a directory\n",
    "output_file_path = \"./fineweb_nemotron_edu_classifier_results\"\n",
    "\n",
    "# Use mode=\"overwrite\" to overwrite the output directory if it already exists\n",
    "# This helps to ensure that the correct output is written\n",
    "write_stage = JsonlWriter(output_file_path, mode=\"overwrite\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initialize Pipeline\n",
    "\n",
    "In NeMo Curator, we use pipelines to run distributed data workflows using Ray. Pipelines take care of resource allocation and autoscaling to achieve enhanced performance and minimize GPU idleness.\n",
    "\n",
    "For the distributed data classifiers, we are able to achieve speedups by ensuring that model inference is run in parallel across all available GPUs, while other stages such as I/O, tokenization, and filtering are run across all available CPUs. This is possible because Curator pipelines are composable, which allows each stage in a pipeline to run independently and with its own specified hardware resources."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Pipeline(name='classifier_pipeline', stages=[jsonl_reader(JsonlReader), nemocurator_fineweb_nemotron_4_edu_classifier_classifier(FineWebNemotronEduClassifier), jsonl_writer(JsonlWriter)])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "classifier_pipeline = Pipeline(name=\"classifier_pipeline\", description=\"Run a classifier pipeline\")\n",
    "\n",
    "# Add stages to the pipeline\n",
    "classifier_pipeline.add_stage(read_stage)\n",
    "classifier_pipeline.add_stage(classifier_stage)\n",
    "classifier_pipeline.add_stage(write_stage)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Composability is also what allows a classifier to sit between pre-processing and post-processing stages. Typical text pre-processing add-ons include text normalization (lowercasing, URL/email removal, Unicode cleanup) and language identification and filtering (to keep only target languages). A full pipeline may look something like:\n",
    "\n",
    "```python\n",
    "pipeline = Pipeline(name=\"full_pipeline\")\n",
    "pipeline.add_stage(read_stage)                # reader (JSONL/S3/etc.)\n",
    "pipeline.add_stage(lang_id_stage)             # optional: language filter\n",
    "pipeline.add_stage(classifier_stage)          # classifier\n",
    "pipeline.add_stage(write_stage)               # writer (JSONL/Parquet)\n",
    "```\n",
    "\n",
    "# Run the  Classifier\n",
    "\n",
    "Let's run the full classifier pipeline:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the pipeline\n",
    "result = classifier_pipeline.run()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since the pipeline ran to completion and the result was written to a JSONL file, we can shut down the Ray cluster with:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    ray_client.stop()\n",
    "except Exception as e:  # noqa: BLE001\n",
    "    print(f\"Error stopping Ray client: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inspect the Output\n",
    "\n",
    "The write stage returns a list of written files. We can read the output file as a Pandas DataFrame for inspection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>text</th>\n",
       "      <th>fineweb-nemotron-edu-score-float</th>\n",
       "      <th>fineweb-nemotron-edu-score-int</th>\n",
       "      <th>fineweb-nemotron-edu-score-label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Get rich quick by following these 5 simple ste...</td>\n",
       "      <td>0.400879</td>\n",
       "      <td>0</td>\n",
       "      <td>low_quality</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Our website uses cookies to improve your exper...</td>\n",
       "      <td>0.239624</td>\n",
       "      <td>0</td>\n",
       "      <td>low_quality</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>The water cycle is a continuous process that m...</td>\n",
       "      <td>3.294922</td>\n",
       "      <td>3</td>\n",
       "      <td>high_quality</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>George Washington was the first president of t...</td>\n",
       "      <td>2.125000</td>\n",
       "      <td>2</td>\n",
       "      <td>low_quality</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Newton's three laws of motion form the foundat...</td>\n",
       "      <td>2.916016</td>\n",
       "      <td>3</td>\n",
       "      <td>high_quality</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                text  \\\n",
       "0  Get rich quick by following these 5 simple ste...   \n",
       "1  Our website uses cookies to improve your exper...   \n",
       "2  The water cycle is a continuous process that m...   \n",
       "3  George Washington was the first president of t...   \n",
       "4  Newton's three laws of motion form the foundat...   \n",
       "\n",
       "   fineweb-nemotron-edu-score-float  fineweb-nemotron-edu-score-int  \\\n",
       "0                          0.400879                               0   \n",
       "1                          0.239624                               0   \n",
       "2                          3.294922                               3   \n",
       "3                          2.125000                               2   \n",
       "4                          2.916016                               3   \n",
       "\n",
       "  fineweb-nemotron-edu-score-label  \n",
       "0                      low_quality  \n",
       "1                      low_quality  \n",
       "2                     high_quality  \n",
       "3                      low_quality  \n",
       "4                     high_quality  "
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# For simplicity, we take the first written file from the writer stage\n",
    "# In real pipelines, the writer may return multiple files (shards) or objects\n",
    "result_file = result[0].data[0]\n",
    "\n",
    "result_df = pd.read_json(result_file, lines=True)\n",
    "result_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that the predictions were generated as expected."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
